# Copyright (c) Alibaba, Inc. and its affiliates.
import os
from dataclasses import dataclass, field
from typing import Literal, Optional

from swift.llm import TEMPLATE_MAPPING
from swift.utils import get_logger

logger = get_logger()


@dataclass
class TemplateArguments:
    """
    TemplateArguments 类是一个数据类，包含与模板配置和使用相关的各种参数。

    参数：
    template (Optional[str]): 模板类型。默认值为 None，表示使用 model_type 对应的模板。
    system (Optional[str]): 覆盖模板中的默认系统。默认值为 None。
    max_length (Optional[int]): 模板的最大长度。默认值为 None。
    truncation_strategy (Literal): 截断模板的策略。默认值为 'delete'。
    max_pixels (Optional[int]): 模板的最大像素数。默认值为 None。
    padding_side: 训练 batch_size >= 2 时的 padding_side。
    loss_scale (str): 训练的损失比例。默认值为 'default'，表示仅计算助手的损失。
    sequence_parallel_size (int): 序列并行的大小。默认值为 1。
    use_chat_template (str): 使用聊天模板或默认生成模板，默认为 True
    template_backend (str): 使用 Swift 模板或 Jinja 模板
    """
    # 模板名称，可选值参考 TEMPLATE_MAPPING。默认值为 None，会自动根据模型选择模板。
    template: Optional[str] = field(
        default=None, metadata={'help': f'template choices: {list(TEMPLATE_MAPPING.keys())}'})

    # 重写模板中默认的 system prompt，例如“你是一个 AI 助手”。可以是字符串或.txt文件路径。
    system: Optional[str] = None

    # 模板的最大 token 长度限制，用于控制模型输入大小。
    max_length: Optional[int] = None

    # 超出最大长度时的截断策略，可选 'delete'（删除溢出部分）、'left'（从左截）、'right'（从右截）。
    truncation_strategy: Literal['delete', 'left', 'right', None] = None

    # 图像输入的最大像素数，通常用于多模态模型（图文任务）。
    max_pixels: Optional[int] = None

    # Agent 专用的模板，比如多轮对话中的代理指令模板。
    agent_template: Optional[str] = None

    # bounding box 归一化策略，用于图像相关任务。'norm1000' 表示归一到 [0,1000]。
    norm_bbox: Literal['norm1000', 'none', None] = None

    # 回复内容的前缀，比如设置为 "答：" 以用于对话标注或 few-shot 模式。
    response_prefix: Optional[str] = None

    # 指定 padding 的方向（左填充或右填充）。训练时通常为 'right'。
    padding_side: Literal['left', 'right'] = 'right'

    # 控制训练时损失计算的策略，'default' 表示仅计算助手回复的损失。
    loss_scale: str = 'default'

    # 序列并行的规模，通常用于大模型的分布式训练，默认值为 1。
    sequence_parallel_size: int = 1

    # 是否使用聊天模板（如对话格式），默认为 True 表示启用聊天格式。
    use_chat_template: bool = True

    # 模板后端选择，可为 'swift'（默认）或 'jinja' 模板引擎。
    template_backend: Literal['swift', 'jinja'] = 'swift'

    def __post_init__(self):
        # 如果 template 没有指定且类中有 model_meta 属性，则自动从模型中获取默认模板
        if self.template is None and hasattr(self, 'model_meta'):
            self.template = self.model_meta.template

        # 如果 system 字段不为空
        if self.system is not None:
            # 如果是文件路径（以 .txt 结尾），则读取文件内容作为 system prompt
            if self.system.endswith('.txt'):
                assert os.path.isfile(self.system), f'self.system: {self.system}'
                with open(self.system, 'r') as f:
                    self.system = f.read()
            else:
                # 否则将字符串中的 \n 替换为换行符
                self.system = self.system.replace('\\n', '\n')

        # 如果定义了 response_prefix，也将其中的 \n 替换为换行符
        if self.response_prefix is not None:
            self.response_prefix = self.response_prefix.replace('\\n', '\n')

        # 如果未设置截断策略，则默认为 'delete'
        if self.truncation_strategy is None:
            self.truncation_strategy = 'delete'

    def get_template_kwargs(self):
        # 将截断策略 'delete' 转换为 'raise'（系统内部处理方式不同）
        truncation_strategy = self.truncation_strategy
        if truncation_strategy == 'delete':
            truncation_strategy = 'raise'

        # 返回一个字典，包含当前模板配置所需的参数，用于调用模板构建逻辑
        return {
            'default_system': self.system,
            'max_length': self.max_length,
            'truncation_strategy': truncation_strategy,
            'max_pixels': self.max_pixels,
            'agent_template': self.agent_template,
            'norm_bbox': self.norm_bbox,
            'response_prefix': self.response_prefix,
            'loss_scale': self.loss_scale,
            'padding_side': self.padding_side,
            'sequence_parallel_size': self.sequence_parallel_size,
            'template_backend': self.template_backend,
            'use_chat_template': self.use_chat_template
        }
